#include "BiCGStab.hpp"

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
BiCGStabSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::BiCGStabSolver() : IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>()
{
}


template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void BiCGStabSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Mult_ipl(const par_Vector<idx_t, ksp_t, dof> & b,
        par_Vector<idx_t, ksp_t, dof> & x) const 
{
    assert(this->oper != nullptr);
    
    MPI_Comm comm = x.comm;
    int my_pid; MPI_Comm_rank(comm, &my_pid);
    double * record = this->part_times;
    memset(record, 0.0, sizeof(double) * NUM_KRYLOV_RECORD);
    ksp_t b_norm, res, r_norm, alpha,
            gamma_numer, gamma_denom, gamma, beta;

    // 初始化辅助向量
    par_Vector<idx_t, ksp_t, dof> r0(x), r(x), s(x), v(x), p(x), q(x);
    par_Vector<idx_t, pc_calc_t, dof> * pc_buf_b = nullptr, * pc_buf_x = nullptr;

    if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
        pc_buf_b = new par_Vector<idx_t, pc_calc_t, dof>(b.comm, b.glb_nrows, b.beg_row, b.end_row);
        pc_buf_x = new par_Vector<idx_t, pc_calc_t, dof>(*pc_buf_b);
        pc_buf_b->set_val(0.0);
        pc_buf_x->set_val(0.0);
    }

    // compute initial residual
    this->oper->Mult(x, r0, false);
    vec_add(b, (ksp_t) -1.0, r0, r0);

    vec_copy(r0, p);
    vec_copy(r0, r);
    b_norm = this->Dot(b, b);    b_norm = sqrt(b_norm);
    res    = this->Dot(r0, r0);  r_norm = sqrt(res);

    if (my_pid == 0) {
        printf("L2 norm of b: %20.16e\n", (double)b_norm);
        printf("Initial L2 norm of residual: %20.16e\n", (double)r_norm);
    }

    int & iter = this->final_iter = 0;
    ksp_t den_norm;
    if (b_norm > 0.0) den_norm = b_norm;// convergence criterion |r_i|/|b| <= accuracy if |b| > 0
    else              den_norm = r_norm;// convergence criterion |r_i|/|r0| <= accuracy if |b| = 0

    /*  convergence criteria: |r_i| <= max( a_tol, r_tol * den_norm)
        den_norm = |r_0| or |b|
        note: default for a_tol is 0.0, so relative residual criteria is used unless
        user specifies a_tol, or sets r_tol = 0.0, which means absolute
        tol only is checked  */
    ksp_t epsilon = std::max(this->abs_tol, this->rel_tol * den_norm);

    // so now our stop criteria is |r_i| <= epsilon
    if (my_pid == 0) {
        if (b_norm > 0.0) {
            printf("=========================================================\n\n");
            printf("Iters             resid.norm               rel.res.norm  \n");
            printf("------       -------------------      --------------------\n");
        } else {
            printf("=============================================\n\n");
            printf("Iters     resid.norm     \n");
            printf("------    ------------    \n");
        }
    }

    /* Start BiCGStab iterations */
    while (iter < this->max_iter)
    {
        iter++;

        if (this->prec) {
            bool zero_guess = true;
            if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                VecTransPrec(p, *pc_buf_b);// copy in
                if (zero_guess) pc_buf_x->set_val(0.0);
                else VecTransPrec(v, *pc_buf_x);// copy in
                
                this->prec->Mult(*pc_buf_b, *pc_buf_x, zero_guess);
                
                VecTransPrec(*pc_buf_x, v);
            } else {
                if (zero_guess) v.set_val(0.0);// (*(bicgstab_functions->ClearVector))(v);
                this->prec->Mult(p, v, zero_guess);// precond(precond_data, A, p, v);
            }
        } else vec_copy(p, v);

        this->oper->Mult(v, q, false);// (*(bicgstab_functions->Matvec))(matvec_data, 1.0, A, v, 0.0, q);
        ksp_t temp = this->Dot(r0, q);// temp = (*(bicgstab_functions->InnerProd))(r0, q);
        alpha = res / temp;

        vec_add(x,  alpha, v, x);// (*(bicgstab_functions->Axpy))(alpha, v, x);
        vec_add(r, -alpha, q, r);// (*(bicgstab_functions->Axpy))(-alpha, q, r);

        if (this->prec) {
            bool zero_guess = true;
            if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                VecTransPrec(r, *pc_buf_b);// copy in
                if (zero_guess) pc_buf_x->set_val(0.0);
                else VecTransPrec(v, *pc_buf_x);// copy in
                
                this->prec->Mult(*pc_buf_b, *pc_buf_x, zero_guess);
                
                VecTransPrec(*pc_buf_x, v);
            } else {
                if (zero_guess) v.set_val(0.0);// (*(bicgstab_functions->ClearVector))(v);
                this->prec->Mult(r, v, zero_guess);// precond(precond_data, A, r, v);
            }
        } else vec_copy(r, v);
      
        this->oper->Mult(v, s, false);// (*(bicgstab_functions->Matvec))(matvec_data, 1.0, A, v, 0.0, s);
        /* Handle case when gamma = 0.0/0.0 as 0.0 and not NAN */
        gamma_numer = this->Dot(r, s);// gamma_numer = (*(bicgstab_functions->InnerProd))(r, s);
        gamma_denom = this->Dot(s, s);// gamma_denom = (*(bicgstab_functions->InnerProd))(s, s);
        if ((gamma_numer == 0.0) && (gamma_denom == 0.0)) gamma = 0.0;
        else                                              gamma = gamma_numer / gamma_denom;
    
        vec_add(x,  gamma, v, x);// (*(bicgstab_functions->Axpy))(gamma, v, x);
        vec_add(r, -gamma, s, r);// (*(bicgstab_functions->Axpy))(-gamma, s, r);
        /* residual is now updated, must immediately check for convergence */
        r_norm = sqrt(this->Dot(r, r));// r_norm = sqrt((*(bicgstab_functions->InnerProd))(r, r));
        if (my_pid == 0) {
            if (b_norm > 0.0)
                printf("%5d    %.16e   %.16e\n", iter, (double)r_norm, (double)(r_norm/b_norm));
            else
                printf("%5d    %.16e\n", iter, (double)r_norm);
        }
        /* check for convergence, evaluate actual residual */
        if (r_norm <= epsilon)
        {
            // vec_copy(b, r);// (*(bicgstab_functions->CopyVector))(b, r);
            // (*(bicgstab_functions->Matvec))(matvec_data, -1.0, A, x, 1.0, r);
            this->oper->Mult(x, r, false);
            vec_add(b, (ksp_t) -1.0, r, r);
            r_norm = this->Norm(r);// r_norm = sqrt((*(bicgstab_functions->InnerProd))(r, r));
            if (r_norm <= epsilon) {
                if (my_pid == 0) printf("Final L2 norm of residual: %e\n\n", r_norm);
                this->converged = 1;
                break;
            }
        }
      
        beta = 1.0 / res;  
        res = this->Dot(r0, r);// res = (*(bicgstab_functions->InnerProd))(r0, r);
        beta *= res;
        vec_add(p, -gamma, q, p);// (*(bicgstab_functions->Axpy))(-gamma, q, p);
        vec_scale(beta * alpha / gamma, p);// (*(bicgstab_functions->ScaleVector))((beta * alpha / gamma), p);
        vec_add(p, (ksp_t) 1.0, r, p);// (*(bicgstab_functions->Axpy))(1.0, r, p);
    } /* end while loop */

    if (pc_buf_x) delete pc_buf_x;
    if (pc_buf_b) delete pc_buf_b;
}

template class BiCGStabSolver<int, double, double, double, 1>;
template class BiCGStabSolver<int, double, double, double, 2>;
template class BiCGStabSolver<int, double, double, double, 3>;
template class BiCGStabSolver<int, double, double, double, 4>;

template class BiCGStabSolver<int, float , double, double, 1>;
template class BiCGStabSolver<int, float , double, double, 2>;
template class BiCGStabSolver<int, float , double, double, 3>;
template class BiCGStabSolver<int, float , double, double, 4>;

template class BiCGStabSolver<int, float , float , double, 1>;
template class BiCGStabSolver<int, float , float , double, 2>;
template class BiCGStabSolver<int, float , float , double, 3>;
template class BiCGStabSolver<int, float , float , double, 4>;

template class BiCGStabSolver<int, float , float , float , 1>;
template class BiCGStabSolver<int, float , float , float , 2>;
template class BiCGStabSolver<int, float , float , float , 3>;
template class BiCGStabSolver<int, float , float , float , 4>;
#ifdef USE_FP16
template class BiCGStabSolver<int, __fp16, float , double, 1>;
template class BiCGStabSolver<int, __fp16, float , double, 2>;
template class BiCGStabSolver<int, __fp16, float , double, 3>;
template class BiCGStabSolver<int, __fp16, float , double, 4>;

template class BiCGStabSolver<int, __fp16, float , float , 1>;
template class BiCGStabSolver<int, __fp16, float , float , 2>;
template class BiCGStabSolver<int, __fp16, float , float , 3>;
template class BiCGStabSolver<int, __fp16, float , float , 4>;
#endif
